"""
改进的数据筛选策略

基于题目要求和轴承故障机理，实现更科学的数据筛选方法
"""

import numpy as np
import pandas as pd
from collections import defaultdict


class ImprovedDataSelector:
    """改进的数据选择器"""
    
    def __init__(self):
        self.selection_criteria = {
            'fault_types': ['B', 'IR', 'OR', 'N'],
            'fault_sizes': ['007', '014', '021', '028'],
            'load_conditions': ['0', '1', '2', '3'],
            'positions': ['@3', '@6', '@12']  # 外圈故障位置
        }
    
    def select_representative_data(self, file_analysis, selection_strategy='systematic'):
        """
        基于轴承故障机理的系统性数据筛选
        
        Args:
            file_analysis: 文件分析结果
            selection_strategy: 选择策略 ('systematic', 'balanced', 'diverse')
        """
        print(f"🎯 执行{selection_strategy}筛选策略...")
        
        if selection_strategy == 'systematic':
            return self._systematic_selection(file_analysis)
        elif selection_strategy == 'balanced':
            return self._balanced_selection(file_analysis)
        elif selection_strategy == 'diverse':
            return self._diverse_selection(file_analysis)
        else:
            raise ValueError(f"不支持的选择策略: {selection_strategy}")
    
    def _systematic_selection(self, file_analysis):
        """
        系统性选择策略 - 基于轴承故障机理
        
        选择原则：
        1. 确保每种故障类型都有代表性样本
        2. 覆盖不同故障尺寸和载荷条件
        3. 考虑外圈故障的不同位置
        4. 优先选择典型工况下的数据
        """
        print("📋 基于轴承故障机理的系统性筛选...")
        
        selected_files = []
        selection_reasoning = {}
        
        # 1. 正常状态数据 - 选择所有可用数据
        normal_files = file_analysis['fault_types']['N']
        selected_files.extend(normal_files)
        selection_reasoning['Normal'] = {
            'count': len(normal_files),
            'reason': '正常状态样本较少，选择全部数据以确保代表性'
        }
        
        # 2. 滚动体故障 - 选择不同故障尺寸和载荷条件的组合
        ball_files = file_analysis['fault_types']['B']
        ball_selected = self._select_ball_fault_samples(ball_files)
        selected_files.extend(ball_selected)
        selection_reasoning['Ball'] = {
            'count': len(ball_selected),
            'reason': '选择不同故障尺寸(007,014,021,028)和载荷条件(0,1,2,3)的典型组合'
        }
        
        # 3. 内圈故障 - 选择不同故障尺寸和载荷条件的组合
        inner_files = file_analysis['fault_types']['IR']
        inner_selected = self._select_inner_fault_samples(inner_files)
        selected_files.extend(inner_selected)
        selection_reasoning['Inner Race'] = {
            'count': len(inner_selected),
            'reason': '选择不同故障尺寸和载荷条件的典型组合，确保内圈故障特征的代表性'
        }
        
        # 4. 外圈故障 - 选择不同位置、故障尺寸和载荷条件的组合
        outer_files = file_analysis['fault_types']['OR']
        outer_selected = self._select_outer_fault_samples(outer_files)
        selected_files.extend(outer_selected)
        selection_reasoning['Outer Race'] = {
            'count': len(outer_selected),
            'reason': '选择不同位置(3点,6点,12点)、故障尺寸和载荷条件的典型组合'
        }
        
        # 打印选择理由
        self._print_selection_reasoning(selection_reasoning)
        
        return selected_files, selection_reasoning
    
    def _select_ball_fault_samples(self, ball_files):
        """选择滚动体故障样本"""
        selected = []
        
        # 按故障尺寸和载荷条件分组
        grouped_files = self._group_files_by_conditions(ball_files)
        
        # 选择策略：每个故障尺寸选择2个不同载荷条件的样本
        for size in ['007', '014', '021', '028']:
            if size in grouped_files:
                size_files = grouped_files[size]
                # 按载荷条件分组
                load_groups = defaultdict(list)
                for file_path in size_files:
                    for load in ['0', '1', '2', '3']:
                        if f'_{load}.mat' in file_path:
                            load_groups[load].append(file_path)
                            break
                
                # 选择不同载荷条件的样本
                selected_loads = ['0', '2']  # 选择0马力和2马力作为典型工况
                for load in selected_loads:
                    if load in load_groups and load_groups[load]:
                        selected.append(load_groups[load][0])  # 选择第一个
                        if len(selected) >= 8:  # 限制总数
                            break
                if len(selected) >= 8:
                    break
        
        return selected[:8]  # 最多选择8个
    
    def _select_inner_fault_samples(self, inner_files):
        """选择内圈故障样本"""
        selected = []
        
        # 按故障尺寸和载荷条件分组
        grouped_files = self._group_files_by_conditions(inner_files)
        
        # 选择策略：每个故障尺寸选择2个不同载荷条件的样本
        for size in ['007', '014', '021', '028']:
            if size in grouped_files:
                size_files = grouped_files[size]
                # 按载荷条件分组
                load_groups = defaultdict(list)
                for file_path in size_files:
                    for load in ['0', '1', '2', '3']:
                        if f'_{load}.mat' in file_path:
                            load_groups[load].append(file_path)
                            break
                
                # 选择不同载荷条件的样本
                selected_loads = ['0', '2']  # 选择0马力和2马力作为典型工况
                for load in selected_loads:
                    if load in load_groups and load_groups[load]:
                        selected.append(load_groups[load][0])  # 选择第一个
                        if len(selected) >= 8:  # 限制总数
                            break
                if len(selected) >= 8:
                    break
        
        return selected[:8]  # 最多选择8个
    
    def _select_outer_fault_samples(self, outer_files):
        """选择外圈故障样本"""
        selected = []
        
        # 按位置分组
        position_groups = {
            '@3': [f for f in outer_files if '@3' in f],
            '@6': [f for f in outer_files if '@6' in f],
            '@12': [f for f in outer_files if '@12' in f]
        }
        
        # 从每个位置选择代表性样本
        for position, files in position_groups.items():
            if not files:
                continue
                
            # 按故障尺寸分组
            grouped_files = self._group_files_by_conditions(files)
            
            # 选择不同故障尺寸的样本
            for size in ['007', '014', '021']:
                if size in grouped_files and grouped_files[size]:
                    # 选择0马力载荷条件的样本
                    for file_path in grouped_files[size]:
                        if '_0.mat' in file_path:
                            selected.append(file_path)
                            break
                    if len(selected) >= 12:  # 限制总数
                        break
                if len(selected) >= 12:
                    break
            if len(selected) >= 12:
                break
        
        return selected[:12]  # 最多选择12个
    
    def _group_files_by_conditions(self, files):
        """按故障尺寸分组文件"""
        grouped = defaultdict(list)
        for file_path in files:
            file_name = file_path.split('\\')[-1]  # 获取文件名
            for size in ['007', '014', '021', '028']:
                if size in file_name:
                    grouped[size].append(file_path)
                    break
        return grouped
    
    def _balanced_selection(self, file_analysis):
        """平衡选择策略 - 确保各类故障样本数量均衡"""
        print("📋 平衡选择策略...")
        
        selected_files = []
        
        # 计算每个故障类型的目标数量
        total_files = sum(len(files) for files in file_analysis['fault_types'].values())
        target_per_type = max(4, total_files // 20)  # 每个类型至少4个，总数不超过20%
        
        for fault_type, files in file_analysis['fault_types'].items():
            if files:
                # 随机选择指定数量的文件
                count = min(target_per_type, len(files))
                selected = np.random.choice(files, count, replace=False)
                selected_files.extend(selected)
                print(f"  {fault_type}: 从{len(files)}个文件中选择了{len(selected)}个")
        
        return selected_files
    
    def _diverse_selection(self, file_analysis):
        """多样性选择策略 - 最大化数据多样性"""
        print("📋 多样性选择策略...")
        
        selected_files = []
        
        # 按故障类型、尺寸、载荷条件、位置进行多样性选择
        all_combinations = []
        
        for fault_type, files in file_analysis['fault_types'].items():
            for file_path in files:
                file_name = file_path.split('\\')[-1]
                
                # 提取特征
                features = {
                    'fault_type': fault_type,
                    'fault_size': self._extract_fault_size(file_name),
                    'load_condition': self._extract_load_condition(file_name),
                    'position': self._extract_position(file_name),
                    'file_path': file_path
                }
                all_combinations.append(features)
        
        # 按多样性排序并选择
        diverse_combinations = self._rank_by_diversity(all_combinations)
        
        # 选择前20个最多样化的组合
        for combo in diverse_combinations[:20]:
            selected_files.append(combo['file_path'])
        
        return selected_files
    
    def _extract_fault_size(self, file_name):
        """提取故障尺寸"""
        for size in ['007', '014', '021', '028']:
            if size in file_name:
                return size
        return 'unknown'
    
    def _extract_load_condition(self, file_name):
        """提取载荷条件"""
        for load in ['0', '1', '2', '3']:
            if f'_{load}.mat' in file_name:
                return load
        return 'unknown'
    
    def _extract_position(self, file_name):
        """提取外圈故障位置"""
        if '@3' in file_name:
            return '@3'
        elif '@6' in file_name:
            return '@6'
        elif '@12' in file_name:
            return '@12'
        return 'none'
    
    def _rank_by_diversity(self, combinations):
        """按多样性排序"""
        # 简单的多样性评分：不同特征组合的数量
        diversity_scores = []
        
        for combo in combinations:
            score = 0
            # 故障类型多样性
            score += len(set([c['fault_type'] for c in combinations if c['fault_type'] == combo['fault_type']]))
            # 故障尺寸多样性
            score += len(set([c['fault_size'] for c in combinations if c['fault_size'] == combo['fault_size']]))
            # 载荷条件多样性
            score += len(set([c['load_condition'] for c in combinations if c['load_condition'] == combo['load_condition']]))
            # 位置多样性
            score += len(set([c['position'] for c in combinations if c['position'] == combo['position']]))
            
            diversity_scores.append((score, combo))
        
        # 按多样性分数排序
        diversity_scores.sort(key=lambda x: x[0], reverse=True)
        
        return [combo for score, combo in diversity_scores]
    
    def _print_selection_reasoning(self, reasoning):
        """打印选择理由"""
        print("\n📋 数据筛选理由说明:")
        print("=" * 50)
        
        for fault_type, info in reasoning.items():
            print(f"\n{fault_type}故障:")
            print(f"  选择数量: {info['count']}个样本")
            print(f"  选择理由: {info['reason']}")
        
        print("\n🎯 总体筛选原则:")
        print("1. 基于轴承故障机理选择典型工况")
        print("2. 确保每种故障类型都有代表性样本")
        print("3. 覆盖不同故障尺寸和载荷条件")
        print("4. 考虑外圈故障的不同位置特征")
        print("5. 优先选择0马力和2马力等典型工况")


